# -*- coding: utf-8 -*-


####Reconstruction of IMERG-Early and SM2RAIN toward the study area

####The IMERG-Early product are available at: https://disc.gsfc.nasa.gov/datasets/GPM_3IMERGDE_06/summary?keywords=early
####The SM2RAIN product are available at: https://zenodo.org/record/6136294#.ZG6uPnZByUl

import numpy as np
import os
import netCDF4 as nc

##IMERG-Early
folder='./IMERG_Early_Daily/'
data_folder=folder+'data/'
save_folder='./IMERG_Early_Daily/'

for year_index in range(2007,2020):
    save_path=save_folder+'CR_0.1/'+str(year_index)+'/'
    interp_save_path=save_folder+'CR_0.01/'+str(year_index)+'/'
    if os.path.isdir(save_path)==False:
        os.mkdir(save_path)
    if os.path.isdir(interp_save_path)==False:
        os.mkdir(interp_save_path)    
    files=os.listdir(data_folder+str(year_index))
    
    for file_index in range(len(files)):
        file_name='CR.'+files[file_index][:-4]
        dataset=nc.Dataset(data_folder+str(year_index)+'/'+files[file_index])
        
        tp = dataset.variables['precipitationCal']
        precipitationCal=tp[:].data
        CR_precipitationCal=precipitationCal[0,2499:3201,1049:1451]
        CR_precipitationCal_inter=np.zeros((7011,4011))

        for i in range(402):
            index=np.arange(0,702)*10
            value=CR_precipitationCal[:,i]
            index_inter=np.arange(0,7011)
            value_inter=np.interp(index_inter, index, value)
            CR_precipitationCal_inter[:,i*10]=value_inter

        for i in range(7011):
            index=np.arange(0,402)*10
            value=CR_precipitationCal_inter[i,index]
            index_inter=np.arange(0,4011)
            value_inter=np.interp(index_inter, index, value)
            CR_precipitationCal_inter[i,:]=value_inter   
        
        np.save(save_path+file_name,CR_precipitationCal.astype(np.float32))
        np.save(interp_save_path+file_name,CR_precipitationCal_inter.astype(np.float32))        

##SM2RAIN
data_folder='./SM2RAIN_2007_2021/'
save_folder='./SM2RAIN_2007_2021/'
filenames=os.listdir(data_folder)

for year_index in range(2007,2020):

    save_path=save_folder+'CR_0.1/'+filenames[year_index-2007][-12:-8]+'/'
    interp_save_path=save_folder+'CR_0.01/'+filenames[year_index-2007][-12:-8]+'/'
    if os.path.isdir(save_path)==False:
        os.mkdir(save_path)
    if os.path.isdir(interp_save_path)==False:
        os.mkdir(interp_save_path)    

    dataset=nc.Dataset(data_folder+filenames[year_index-2007]) 
    tp = dataset.variables['Rainfall']
    prec=tp[:].data
        
    CR_data=prec[1049:1452,2499:3202,:]  
    for day in range(CR_data.shape[2]):
        if day<10:
            day_id='00'+str(day)
        elif day<100:
            day_id='0'+str(day)
        else:
            day_id=day
        file_name=filenames[year_index-2007][:-2]+str(day_id)
        CR_data_day=CR_data[:,:,day].T
        CR_data_day_inter=np.zeros((7021,4021))

        for i in range(403):
            index=np.arange(0,703)*10
            value=CR_data_day[:,i]
            index_inter=np.arange(0,7021)
            value_inter=np.interp(index_inter, index, value)
            CR_data_day_inter[:,i*10]=value_inter
        
        for i in range(7021):
            index=np.arange(0,403)*10
            value=CR_data_day_inter[i,index]
            index_inter=np.arange(0,4021)
            value_inter=np.interp(index_inter, index, value)
            CR_data_day_inter[i,:]=value_inter 
        
        Final_CR_data_day_inter=CR_data_day_inter[5:-5,5:-5]
        Final_CR_data_day=np.zeros((702,402))
        lat_index=np.arange(0,402)*10
        lon_index=np.arange(0,702)*10
        for  i in range(702):
            Final_CR_data_day[i,:]=Final_CR_data_day_inter[i*10,lat_index]
            
        np.save(save_path+file_name,Final_CR_data_day.astype(np.float32))
        np.save(interp_save_path+file_name,Final_CR_data_day_inter.astype(np.float32))        

